#!/usr/bin/env python3

import json
import ssl
from typing import Optional

from fastapi import FastAPI, WebSocket, WebSocketDisconnect
from fastapi.responses import JSONResponse, PlainTextResponse

app = FastAPI(title="WebSocket Echo Server", version="1.0.0")


@app.get("/")
async def root():
    return JSONResponse({"message": "WebSocket echo server is running", "ws_path": "/ws"})


@app.get("/healthz")
async def healthz():
    return PlainTextResponse("ok")


@app.websocket("/ws")
async def websocket_endpoint(ws: WebSocket):
    # 可在此进行简单的鉴权，例如校验 query 参数 token
    # token = ws.query_params.get("token")
    # if token != "expected":
    #     await ws.close(code=4401)
    #     return

    await ws.accept()
    try:
        while True:
            message = await ws.receive()
            # message 可能包含 'text' 或 'bytes'
            if "text" in message and message["text"] is not None:
                text = message["text"]
                # 简单 echo：原样返回文本
                await ws.send_text(text)
            elif "bytes" in message and message["bytes"] is not None:
                data = message["bytes"]
                await ws.send_bytes(data)
            elif message.get("type") == "websocket.disconnect":
                break
    except WebSocketDisconnect:
        # 客户端断开
        pass
    except Exception as e:
        # 发生异常时尝试关闭连接
        try:
            await ws.close(code=1011)
        except Exception:
            pass
        raise e


# 兼容：允许直接连接根路径 '/' 进行 WebSocket，避免误用路径导致 403
@app.websocket("/")
async def websocket_root(ws: WebSocket):
    await websocket_endpoint(ws)


if __name__ == "__main__":
    import uvicorn
    
    # SSL 配置示例（需要证书文件）
    # ssl_context = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER)
    # ssl_context.load_cert_chain("cert.pem", "key.pem")
    
    # 启动服务器
    # HTTP (ws://)
    uvicorn.run(app, host="0.0.0.0", port=8000)
    
    # HTTPS (wss://) - 需要证书文件
    # uvicorn.run(app, host="0.0.0.0", port=8443, ssl_context=ssl_context)
